from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains.history_aware_retriever import create_history_aware_retriever
from langchain.chains.retrieval import create_retrieval_chain
from langchain_community.vectorstores import Chroma
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableWithMessageHistory
from langchain_community.chat_message_histories import ChatMessageHistory

from Embed_model import BGE_Embed
from qw_model import qw_model

# 千问大模型
model = qw_model()

embedding = BGE_Embed()

# 加载知识库
vectorstore = Chroma(
    embedding_function=embedding,
    persist_directory='./chroma_rag'
)

# 创建检索器
retriever = vectorstore.as_retriever(search_kwargs={'k': 3})

# 检索出的文档
# result = retriever.invoke('人是你们公司负责还是我们自己找')
# print(result)
# exit()

contextualize_q_system_prompt = (
    "给定聊天记录和最新用户问题 "
    "这可能会引用聊天历史中的上下文， "
    "制定一个可以理解的独立问题 "
    "没有聊天记录。不要回答问题， "
    "如果需要，只需重新规划，否则按原样退回。"
)
contextualize_q_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", contextualize_q_system_prompt),
        MessagesPlaceholder("chat_history"),
        ("human", "{input}"),
    ]
)
history_aware_retriever = create_history_aware_retriever(
    model, retriever, contextualize_q_prompt
)

system_prompt = (
    " 你是一个聪明又有亲和力的销售客服，名叫银弓，擅长用轻松、亲切的口语化的语气解答用户关于公司业务、价格、产品咨询的问题。你的目标是提供专业、准确的回答，尤其是价格问题，严格按照相应文档回答，风格需要口语化、自然流畅、贴近用户生活，避免书面化和过于正式的表达以及尊称。"
    "1. 当用户问题不清楚时，主动询问用户具体需求。"
    "2. 回答价格问题，严格按照检索的文档给出，不可自行编造和总结价格范围。"
    "3. 当没有用户所说的地区（ip），就按它的省份价格算，当省份也没有那就只能做不限地区的"
    "4. 保持沟通语气亲切、自然，避免过于正式的措辞，比如“老板”、“宝子”、“咱们”，不要使用'您'，'米'可以代表'钱'的意思。"
    "5. 输出答案应简洁明了，字数控制在15字左右。"
    "\n\n"
    "{context}"
)

# 创建提示词模板
qa_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system_prompt),
        MessagesPlaceholder("chat_history"),
        ("human", "{input}"),
    ]
)

question_answer_chain = create_stuff_documents_chain(model, qa_prompt)

rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)

# 有状态地管理聊天记录
store = {}
def get_session_history(session_id: str) -> BaseChatMessageHistory:
    if session_id not in store:
        store[session_id] = ChatMessageHistory()
    return store[session_id]


Final_chain = RunnableWithMessageHistory(
    rag_chain,
    get_session_history,
    input_messages_key="input",
    history_messages_key="chat_history",
    output_messages_key="answer",
)

resp1 = Final_chain.invoke({"input": "真人口播多少钱"},
                           config={
                               "configurable": {"session_id": "abc123"}
                           })["answer"]
print(resp1)
#
resp3 = Final_chain.invoke({"input": "贵阳，要三级的"},
                           config={
                               "configurable": {"session_id": "abc123"}
                           })["answer"]
print(resp3)
resp4 = Final_chain.invoke({"input": "人是你们公司负责还是我们自己找"},
                           config={
                               "configurable": {"session_id": "abc123"}
                           })["answer"]
print(resp4)
#
# resp2 = Final_chain.invoke({"input": "那价格是多少"},
#                            config={
#                                "configurable": {"session_id": "abc123"}
#                            })["answer"]
# print(resp2)


# print(store["abc123"].messages)



#保存消息
# dics = messages_to_dict(store["abc123"].messages)
#
#
# import pickle
# f = open("./memory", "wb")
# pickle.dump(dics, f)
# f.close()